import math
import os
import numpy as np
import cv2
import re
from concurrent.futures import ThreadPoolExecutor


def load_image(image_dir, row, col):
    # 正则表达式匹配前缀部分和tr-row-tc-col部分
    pattern = re.compile(r'(.*)tr(\d+)-tc(\d+)\.png')

    # 遍历文件夹中的所有文件
    for filename in os.listdir(image_dir):
        # 检查文件名是否符合完整的tr{row}-tc{col}格式
        match = pattern.match(filename)
        if match:
            prefix = match.group(1)
            matched_row = int(match.group(2))
            matched_col = int(match.group(3))
            if matched_row == row and matched_col == col:
                filepath = os.path.join(image_dir, filename)
                image = cv2.imread(filepath)
                if image is not None:
                    return image, prefix, filename
                else:
                    print(f"Error reading image at {filepath}")

    # 如果没有匹配的文件，返回None
    return None, None, None


def save_image(image, filename, output_dir):
    filepath = os.path.join(output_dir, filename)
    cv2.imwrite(filepath, image)
    print(f"Saved {filename} to {filepath}")


def process_subgrid(subgrid_row, subgrid_col, rows, cols, image_dir, image_output_dir,
                    image_tile_size, image_subgrid_rows, image_subgrid_cols, row_overlap, col_overlap,
                    prefix):
    subgrid_id = (subgrid_row, subgrid_col)
    image_combined = np.zeros((image_subgrid_rows * image_tile_size, image_subgrid_cols * image_tile_size, 3),
                              dtype=np.uint8)

    for i in range(subgrid_row, subgrid_row + image_subgrid_rows):
        for j in range(subgrid_col, subgrid_col + image_subgrid_cols):
            image_part, temp_prefix, filename = load_image(image_dir, i, j)
            row_offset = (i - subgrid_row) * image_tile_size
            col_offset = (j - subgrid_col) * image_tile_size
            image_combined[row_offset:row_offset + image_tile_size, col_offset:col_offset + image_tile_size,
            :] = image_part

            if prefix is None:
                prefix = temp_prefix

            print(
                f"Image {filename} is merged at position ({i}, {j}) into the combined image at offset ({row_offset}, {col_offset}).")
    save_image(image_combined,
               f"{prefix}_tr{math.ceil((subgrid_row - 1) / (image_subgrid_rows - row_overlap)) + 1}-tc{math.ceil((subgrid_col - 1) / (image_subgrid_cols - col_overlap)) + 1}.png",
               image_output_dir)


def combine_images(image_dir, image_output_dir):
    rows, cols = 37, 33
    image_tile_size = 4096
    image_subgrid_rows, image_subgrid_cols = 7, 7
    row_overlap, col_overlap = 2, 2

    os.makedirs(image_output_dir, exist_ok=True)

    prefix = None
    processed_subgrids = set()
    subgrid_tasks = []

    with ThreadPoolExecutor(max_workers=8) as executor:
        for subgrid_row in range(1, rows + 1, image_subgrid_rows - row_overlap):
            if subgrid_row + image_subgrid_rows - 1 > rows:
                subgrid_row = rows - image_subgrid_rows + 1
            for subgrid_col in range(1, cols + 1, image_subgrid_cols - col_overlap):
                if subgrid_col + image_subgrid_cols - 1 > cols:
                    subgrid_col = cols - image_subgrid_cols + 1

                subgrid_id = (subgrid_row, subgrid_col)
                if subgrid_id in processed_subgrids:
                    continue
                processed_subgrids.add(subgrid_id)

                task = executor.submit(process_subgrid, subgrid_row, subgrid_col, rows, cols, image_dir,
                                       image_output_dir, image_tile_size, image_subgrid_rows, image_subgrid_cols,
                                       row_overlap, col_overlap, prefix)
                subgrid_tasks.append(task)

        for task in subgrid_tasks:
            task.result()


# 示例用法
image_dir = '/root/data/deep_align/11866'  # 替换为实际图像文件夹路径
image_output_dir = '/root/data/deep_align/deep_align/11866_big'  # 替换为实际图像合并输出文件夹路径

combine_images(image_dir, image_output_dir)
